from torch import nn

from ssd.modeling.backbone import build_backbone
from ssd.modeling.box_head import build_box_head


class SSDDetector(nn.Module):
    '''
    ssd 目标检测网络
    '''
    def __init__(self, cfg):
        '''
        cfg: 配置文件
        todo 看到这里
        '''
        super().__init__()
        self.cfg = cfg
        # 创建骨干网络，用于图片特征的提取
        self.backbone = build_backbone(cfg)
        # 创建
        self.box_head = build_box_head(cfg)

    def forward(self, images, targets=None):
        features = self.backbone(images)
        detections, detector_losses = self.box_head(features, targets)
        if self.training:
            return detector_losses
        return detections
